import gradio as gr
import os
import json
import time
import shutil
import uuid

from Core.comfy_client import ComfyUIClient


class CannyReduxInterface:
    def __init__(self):
        self.client = ComfyUIClient()
        # 目标工作流：flux_canny_redux.json
        self.workflow_path = os.path.join(
            os.path.dirname(__file__), "..", "Comfyui", "work_flow", "flux_canny_redux.json"
        )

        # 目录：loras、input、temp
        from config.config import get_lora_dir_path, get_input_path
        self.lora_dir = get_lora_dir_path()
        self.input_dir = get_input_path()
        self.temp_dir = os.path.join(os.path.dirname(__file__), "..", "temp")
        os.makedirs(self.temp_dir, exist_ok=True)
        os.makedirs(self.input_dir, exist_ok=True)

        # 收集 LoRA 选项（包含子目录）
        self.lora_choices = self._collect_lora_choices()
        self.default_lora = "Flux\\ETShoeF.safetensors"
        if self.lora_choices and self.default_lora not in self.lora_choices:
            self.default_lora = self.lora_choices[0]

    def _collect_lora_choices(self):
        choices = []
        try:
            if os.path.exists(self.lora_dir):
                for root, _, files in os.walk(self.lora_dir):
                    for f in files:
                        if f.lower().endswith(('.safetensors', '.ckpt')):
                            rel = os.path.relpath(os.path.join(root, f), self.lora_dir)
                            # 使用反斜杠，保持与工作流示例一致
                            rel = rel.replace('/', '\\')
                            choices.append(rel)
        except Exception as e:
            print(f"收集LoRA失败: {e}")
        choices.sort()
        return choices

    def load_workflow(self):
        try:
            with open(self.workflow_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"加载工作流失败: {e}")
            return None

    def _copy_image_to_input(self, image_path, prefix):
        """复制图片到 ComfyUI input，返回文件名"""
        try:
            if not image_path or not os.path.exists(image_path):
                return None, "错误：未提供有效图片路径"
            ext = os.path.splitext(image_path)[1].lower()
            if ext not in [".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff"]:
                return None, "错误：仅支持PNG/JPG/JPEG/BMP/WEBP/TIFF"
            filename = f"{prefix}_{uuid.uuid4().hex}{ext}"
            target = os.path.join(self.input_dir, filename)
            shutil.copy2(image_path, target)
            return filename, None
        except Exception as e:
            return None, f"复制图片失败: {e}"

    def generate_image(
        self,
        lora_name,
        strength_model,
        strength_clip,
        style_strength,
        ref_image_path,
        struct_image_path,
        low_threshold,
        high_threshold,
        ctrl_strength,
        start_percent,
        end_percent,
        steps,
        denoise,
        max_retries: int = 40,
    ):
        try:
            workflow = self.load_workflow()
            if not workflow:
                return None, None, "生成失败：无法加载工作流"

            # 节点81：LoRA设置
            if "81" in workflow and "inputs" in workflow["81"]:
                workflow["81"]["inputs"]["lora_name"] = lora_name
                workflow["81"]["inputs"]["strength_model"] = float(strength_model)
                workflow["81"]["inputs"]["strength_clip"] = float(strength_clip)

            # 节点73：StyleModelApply 强度
            if "73" in workflow and "inputs" in workflow["73"]:
                workflow["73"]["inputs"]["strength"] = float(style_strength)

            # 复制参考图（节点52）
            ref_filename, err1 = self._copy_image_to_input(ref_image_path, "canny_ref")
            if err1:
                return None, None, f"生成失败：{err1}"
            if "52" in workflow and "inputs" in workflow["52"]:
                workflow["52"]["inputs"]["image"] = ref_filename

            # 复制结构图（节点66）
            struct_filename, err2 = self._copy_image_to_input(struct_image_path, "canny_struct")
            if err2:
                return None, None, f"生成失败：{err2}"
            if "66" in workflow and "inputs" in workflow["66"]:
                workflow["66"]["inputs"]["image"] = struct_filename

            # Canny（节点65）阈值
            if "65" in workflow and "inputs" in workflow["65"]:
                workflow["65"]["inputs"]["low_threshold"] = float(low_threshold)
                workflow["65"]["inputs"]["high_threshold"] = float(high_threshold)

            # ControlNetApplyAdvanced（节点72）强度/开始/结束
            if "72" in workflow and "inputs" in workflow["72"]:
                workflow["72"]["inputs"]["strength"] = float(ctrl_strength)
                workflow["72"]["inputs"]["start_percent"] = float(start_percent)
                workflow["72"]["inputs"]["end_percent"] = float(end_percent)

            # KSampler（节点83）步数/降噪
            if "83" in workflow and "inputs" in workflow["83"]:
                workflow["83"]["inputs"]["steps"] = int(steps)
                workflow["83"]["inputs"]["denoise"] = float(denoise)

            # 提交
            resp = self.client.post_prompt(workflow)
            prompt_id = resp.get("prompt_id")
            if not prompt_id:
                return None, None, "生成失败：未返回prompt_id"

            # 轮询，抓取结果（节点59）与canny预览（节点64）
            retry_interval = 2
            preview_path = None
            for _ in range(max_retries):
                history = self.client.get_history(prompt_id)
                if prompt_id in history and "outputs" in history[prompt_id]:
                    outputs = history[prompt_id]["outputs"]

                    # 提取canny预览（节点64）
                    try:
                        node64 = outputs.get("64")
                        if node64 and "images" in node64 and node64["images"]:
                            img = node64["images"][0]
                            img_bytes = self.client.get_image(img["filename"], img.get("subfolder"), img.get("type"))
                            preview_path = os.path.join(self.temp_dir, f"canny_preview_{prompt_id}.png")
                            with open(preview_path, "wb") as f:
                                f.write(img_bytes)
                    except Exception:
                        pass

                    # 结果图：优先节点59（VAEDecode）
                    node59 = outputs.get("59")
                    if node59 and "images" in node59 and node59["images"]:
                        image = node59["images"][0]
                        img_bytes = self.client.get_image(image["filename"], image.get("subfolder"), image.get("type"))
                        result_path = os.path.join(self.temp_dir, f"final_{prompt_id}.png")
                        with open(result_path, "wb") as f:
                            f.write(img_bytes)
                        # 复制到项目 output_images
                        try:
                            from config.config import get_local_output_images_path
                            out_dir = get_local_output_images_path()
                            os.makedirs(out_dir, exist_ok=True)
                            shutil.copy2(result_path, os.path.join(out_dir, os.path.basename(result_path)))
                        except Exception as e:
                            print(f"复制到output_images失败: {e}")
                        return os.path.abspath(result_path), os.path.abspath(preview_path) if preview_path else None, "生成成功"

                    # 备选：节点62（SaveImage）的输出
                    node62 = outputs.get("62")
                    if node62 and "images" in node62 and node62["images"]:
                        image = node62["images"][0]
                        img_bytes = self.client.get_image(image["filename"], image.get("subfolder"), image.get("type"))
                        result_path = os.path.join(self.temp_dir, f"final_{prompt_id}.png")
                        with open(result_path, "wb") as f:
                            f.write(img_bytes)
                        try:
                            from config.config import get_local_output_images_path
                            out_dir = get_local_output_images_path()
                            os.makedirs(out_dir, exist_ok=True)
                            shutil.copy2(result_path, os.path.join(out_dir, os.path.basename(result_path)))
                        except Exception as e:
                            print(f"复制到output_images失败: {e}")
                        return os.path.abspath(result_path), os.path.abspath(preview_path) if preview_path else None, "生成成功"

                time.sleep(retry_interval)

            return None, preview_path, "生成失败：超时或无输出"
        except Exception as e:
            print(f"生成异常: {e}")
            return None, None, f"生成失败：{e}"


def Tab_canny_redux():
    interface = CannyReduxInterface()

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 根据结构参考图设计")

            # 节点81：LoRA选择 + 强度设置
            lora_dropdown = gr.Dropdown(
                label="LoRA名称(lora_name)",
                choices=interface.lora_choices,
                value=interface.default_lora,
                interactive=bool(interface.lora_choices),
            )
            strength_model_slider = gr.Slider(label="模型强度(strength_model)", minimum=0, maximum=1, value=1, step=0.01)
            strength_clip_slider = gr.Slider(label="CLIP强度(strength_clip)", minimum=0, maximum=1, value=1, step=0.01)

            # 节点52：加载参考图
            ref_image = gr.Image(label="加载参考图", type="filepath")

            # 节点73：样式强度
            style_strength_slider = gr.Slider(label="风格强度(节点73 strength)", minimum=0, maximum=1, value=1, step=0.01)

            # 节点66：加载结构图
            struct_image = gr.Image(label="加载结构图", type="filepath")

            # 节点65：Canny 阈值
            low_threshold_slider = gr.Slider(label="Canny低阈值", minimum=0, maximum=1, value=0.15, step=0.01)
            high_threshold_slider = gr.Slider(label="Canny高阈值", minimum=0, maximum=1, value=0.50, step=0.01)

            # 节点72：ControlNet应用(高级)
            ctrl_strength_slider = gr.Slider(label="ControlNet强度", minimum=0, maximum=1, value=1, step=0.01)
            start_percent_slider = gr.Slider(label="开始时间(百分比)", minimum=0, maximum=1, value=0, step=0.01)
            end_percent_slider = gr.Slider(label="结束时间(百分比)", minimum=0, maximum=1, value=0.4, step=0.01)

            # 节点83：KSampler
            steps_input = gr.Number(label="步数(steps)", value=20, minimum=0, maximum=60)
            denoise_slider = gr.Slider(label="降噪(denoise)", minimum=0, maximum=1, value=1, step=0.01)

        with gr.Column(scale=2):
            generate_btn = gr.Button("生成", variant="primary")
            result_image = gr.Image(label="生成结果")
            status_text = gr.Textbox(label="状态", interactive=False)
            canny_preview = gr.Image(label="预览canny线条")
            # 输出文件夹图片预览
            from config.config import get_output_path
            output_dir = get_output_path()
            def get_output_images():
                if not os.path.exists(output_dir):
                    return []
                images = []
                for file in os.listdir(output_dir):
                    if file.lower().endswith((".png", ".jpg", ".jpeg")):
                        images.append(os.path.join(output_dir, file))
                images.sort(key=lambda x: os.path.getmtime(x), reverse=True)
                return images
            with gr.Row():
                refresh_btn = gr.Button("刷新输出预览")
            with gr.Row():
                output_gallery = gr.Gallery(label="输出文件夹图片", columns=4, show_label=True, elem_id="output_gallery_i2i")

            refresh_btn.click(fn=get_output_images, outputs=[output_gallery])


    def run_generate(lora_name, strength_model, strength_clip, style_strength, ref_image_path, struct_image_path,
                     low_t, high_t, ctrl_s, start_p, end_p, steps, denoise):
        res_path, preview_path, status = interface.generate_image(
            lora_name, strength_model, strength_clip, style_strength,
            ref_image_path, struct_image_path,
            low_t, high_t, ctrl_s, start_p, end_p,
            steps, denoise
        )
        return res_path, status, preview_path

    generate_btn.click(
        fn=run_generate,
        inputs=[
            lora_dropdown,
            strength_model_slider,
            strength_clip_slider,
            style_strength_slider,
            ref_image,
            struct_image,
            low_threshold_slider,
            high_threshold_slider,
            ctrl_strength_slider,
            start_percent_slider,
            end_percent_slider,
            steps_input,
            denoise_slider,
        ],
        outputs=[result_image, status_text, canny_preview],
    )

    
  